from common_import import *


def calculate_focus_stats(data, focus):
    # 提取商品代码和数量
    codes = data["code"]
    quantities = data["quantity"]

    # 创建一个列表来存储结果
    result_list = []

    for code in focus:
        # 过滤目标商品的记录
        code_mask = codes == code

        # 计算购买总次数和总量
        purchase_count = code_mask.sum()
        total_quantity = quantities[code_mask].sum()

        # 将结果添加到结果列表中
        result_list.append((code, purchase_count, total_quantity))

    # 将结果列表转换为结构化数组
    result_array = np.array(
        result_list,
        dtype=[("code", "i4"), ("purchase_count", "i4"), ("total_quantity", "f4")],
    )

    return result_array


def calculate_entropy_topsis(data):
    # 提取数据
    purchase_counts = data["purchase_count"]
    total_quantities = data["total_quantity"]

    # 构建决策矩阵
    X = np.vstack((purchase_counts, total_quantities)).T

    # 标准化决策矩阵
    R = X / np.sqrt(np.sum(X**2, axis=0))

    # 计算熵值
    P = R / np.sum(R, axis=0)
    entropy = -np.sum(P * np.log(P + 1e-9), axis=0) / np.log(len(X))

    # 计算权重
    weights = (1 - entropy) / np.sum(1 - entropy)

    # 加权标准化决策矩阵
    Z = R * weights

    # 确定正理想解和负理想解
    Z_max = Z.max(axis=0)
    Z_min = Z.min(axis=0)

    # 计算与正理想解和负理想解的距离
    S_plus = np.sqrt(np.sum((Z - Z_max) ** 2, axis=1))
    S_minus = np.sqrt(np.sum((Z - Z_min) ** 2, axis=1))

    # 计算相对接近度
    C = S_minus / (S_plus + S_minus)

    # 将相对接近度作为新列添加到原始结构化数组
    result_array = np.zeros(
        data.shape,
        dtype=[
            ("code", "i4"),
            ("purchase_count", "i4"),
            ("total_quantity", "f4"),
            ("demand_index", "f4"),
        ],
    )
    result_array["code"] = data["code"]
    result_array["purchase_count"] = data["purchase_count"]
    result_array["total_quantity"] = data["total_quantity"]
    result_array["demand_index"] = C

    return result_array


def plot_top5_demand_index(result):
    # # 按需求指数排序，取前五个商品
    # sorted_indices = np.argsort(data["demand_index"])[::-1][:5]  # 按需求指数降序排序
    # top5_codes = data["code"][sorted_indices]
    # top5_demand_indices = data["demand_index"][sorted_indices]

    # # 创建柱状图
    # plt.figure(figsize=(10, 6))
    # plt.bar(top5_codes, top5_demand_indices, color="skyblue")
    # plt.xlabel("Product Code")
    # plt.ylabel("Demand Index")
    # plt.title("Top 5 Products by Demand Index")
    # plt.show()
    codes = result["code"]
    purchase_counts = result["purchase_count"]
    total_quantities = result["total_quantity"]
    demand_indices = result["demand_index"]
    # 1. 柱状图 - 显示市场需求指数
    plt.figure(figsize=(10, 6))
    plt.bar(codes, demand_indices, color="skyblue")
    plt.xlabel("产品编号")
    plt.ylabel("市场需求指数")
    plt.title("初筛得到产品的市场需求指数")
    mydraw.show_or_print("demand_index.png")


if __name__ == "__main__":
    pass
    data = tool.get_np("focus_product_stat.csv")
    plot_top5_demand_index(data)
